from core import Logger, Config
from core.utils.pool import BrowserPool
from celery import signals, Celery
from playwright.sync_api import Page, Route
import re
pool: BrowserPool = None
         
dlApp = Celery(
    "downloader",
    broker=f'redis://:{Config.REDIS_PASSWORD}@{Config.REDIS_HOST}:{Config.REDIS_PORT}/{Config.REDIS_DB}',
    backend=f'redis://:{Config.REDIS_PASSWORD}@{Config.REDIS_HOST}:{Config.REDIS_PORT}/{Config.REDIS_DB}'
)

dlApp.conf.broker_connection_retry_on_startup = True

@dlApp.task(bind=True)
def download(self, url: str, **kwargs) ->str:
    """
    异步下载页面内容。
    
    参数:
    - url: 要下载的页面的URL，类型为字符串。
    - **kwargs: 传递给下载器的额外关键字参数。
    
    返回值:
    - 下载的页面内容，类型为字符串。
    """
    page: Page = pool.acquire_page()
    def handleRoute(route:Route, request):
        route.abort()

    page.route(re.compile(r"\.(jpg|png|svg)$"), handleRoute)
    page.goto(url)
    page.wait_for_load_state('networkidle')
    content: str = page.content()
    page.unroute_all()
    pool.release_page(page)
    return content


# 定义初始化函数  
@signals.worker_process_init.connect
def init_worker(**kwargs):  
    Logger.debug("Initializing worker...")  
    global pool
    pool = BrowserPool(headless=Config.POOL_HEADLESS, 
                        workdir=Config.POOL_WORK_DIR,
                        pool_size=Config.POOL_SIZE)
    pool.initialize()

@signals.worker_process_shutdown.connect
def close_worker(**kwargs):  
    Logger.debug("Closing worker...")  
    global pool
    pool.close()

# main thread
#signals.worker_init.connect(init_worker)
#signals.worker_shutdown.connect(close_worker)